package com.gizwits.util

import java.util.Properties
import kafka.common.KafkaException
import kafka.consumer.ConsumerConfig
import kafka.utils.{ZKStringSerializer, ZKGroupTopicDirs, ZkUtils}
import org.I0Itec.zkclient.ZkClient

object UpdateOffsetsInZK {
  def createConsumerConfig: ConsumerConfig = {
    val props = new Properties()
    props.put("zookeeper.connect", "s60:2181")
    props.put("group.id", "1")
    props.put("auto.offset.reset", "smallest")
    props.put("zookeeper.session.timeout.ms", "400")
    props.put("zookeeper.sync.time.ms", "200")
    props.put("auto.commit.interval.ms", "10")
    val config = new ConsumerConfig(props)
    config
  }
  private def getAndSetOffsets(zkClient: ZkClient, config: ConsumerConfig, topic: String,offset:String): Unit = {
    val partitionsPerTopicMap = ZkUtils.getPartitionsForTopics(zkClient, List(topic))
    var partitions: Seq[Int] = Nil

    //topic
    partitionsPerTopicMap.get(topic) match {
      case Some(l) =>  partitions = l.sortWith((s,t) => s < t)
      case _ => throw new RuntimeException("Can't find topic " + topic)
    }

    var numParts = 0
    for (partition <- partitions) {
      //  brokerhost
      val brokerHostingPartition = ZkUtils.getLeaderForPartition(zkClient, topic, partition)

      val broker = brokerHostingPartition match {
        case Some(b) => b
        case None => throw new KafkaException("Broker " + brokerHostingPartition + " is unavailable. Cannot issue " +
          "getOffsetsBefore request")
      }
      ZkUtils.getBrokerInfo(zkClient, broker) match {

        //使用的 低级api  获取zk 信息 ,并写入offset 到zk
        case Some(brokerInfo) =>
          val topicDirs = new ZKGroupTopicDirs(config.groupId, topic)
         val zkdata= ZkUtils.readData(zkClient, topicDirs.consumerOffsetDir + "/" + partition)
          println(zkdata._1+"--"+zkdata._2)
          ZkUtils.updatePersistentPath(zkClient, topicDirs.consumerOffsetDir + "/" + partition, offset)

          numParts += 1
        case None => throw new KafkaException("Broker information for broker id %d does not exist in ZK".format(broker))
      }
    }
    println("updated the offset for " + numParts + " partitions")
  }
  def zkOffset(topic:String,offset:String)={
    val zkClient = new ZkClient(createConsumerConfig.zkConnect, createConsumerConfig.zkSessionTimeoutMs,
      createConsumerConfig.zkConnectionTimeoutMs, ZKStringSerializer)
    getAndSetOffsets(zkClient,createConsumerConfig,topic,offset)
  }
  def main(args: Array[String]) {
    zkOffset("gizwits_raw",11068.toString)
  }
}
